import tensorflow as tf
import io
import matplotlib.pyplot as plt
from sklearn import metrics

class ConfusionMatrixMetrics(tf.keras.metrics.Metric):

  def __init__(self, num_class,name='confusion_matrix', **kwargs):
    super(ConfusionMatrixMetrics, self).__init__(name=name, **kwargs)
    self.class_num = num_class
    self.true_positives = self.add_weight(name='matrix', initializer='zeros',shape=(num_class,num_class))

  def update_state(self, y_true, y_pred, sample_weight=None):
    y_true = tf.cast(y_true, tf.float32)
    y_true = tf.expand_dims(y_true,axis=1)
    # y_pred = tf.cast(y_pred, tf.bool)
    y_pred = tf.one_hot(tf.argmax(y_pred,axis=-1),self.class_num,dtype=tf.float32)
    y_pred = tf.expand_dims(y_pred,axis=2)
    res = y_true*y_pred
    self.true_positives.assign_add(tf.reduce_sum(res,axis=0))

  def result(self):
    return self.true_positives

  def reset_states(self):
    self.true_positives.assign(tf.zeros(shape=(self.class_num,self.class_num)))

def _plot_to_image(figure):
  buf = io.BytesIO()
  plt.savefig(buf, format='png')
  plt.close(figure)
  buf.seek(0)
  # Convert PNG buffer to TF image
  image = tf.image.decode_png(buf.getvalue(), channels=4)
  # Add the batch dimension
  image = tf.expand_dims(image, 0)
  return image

def confusion_matrix(summary_writer, matrix, class_names, epoch):
  cm = matrix
  figure = metrics.ConfusionMatrixDisplay(confusion_matrix=cm,display_labels=class_names).plot().figure_
  cm_image = _plot_to_image(figure)
  with summary_writer.as_default():
    tf.summary.image("Confusion Matrix", cm_image, step=epoch)

# if __name__ == '__main__':
#     m = ConfusionMatrixMetrics(4)
#     m.update_state([[0,1,0,0],[0,0,0,1],[0,1,0,0],[0,0,0,1]], [[0.1,0.1,0.2,0.7],[0.2,0.2,0.1,0.5],[0.1,0.4,0.2,0.3],[0.2,0.2,0.1,0.5]])
#     confusion_matrix(None,m.result(),["a","b","c","d"],0)
